// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/row_number_fuzzer.h>

#include <utility>
#include <pollux/common/file/file_systems.h>
#include <pollux/testing/exec/fuzzer_util.h>
#include <pollux/testing/exec/row_number_fuzzer_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/common/file/temp_directory_path.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec {
    namespace {
        class RowNumberFuzzer : public RowNumberFuzzerBase {
        public:
            explicit RowNumberFuzzer(
                size_t initialSeed,
                std::unique_ptr<ReferenceQueryRunner>);

        private:
            // Runs one test iteration from query plans generations, executions and result
            // verifications.
            void runSingleIteration() override;

            std::pair<std::vector<std::string>, std::vector<TypePtr> >
            generatePartitionKeys();

            std::vector<RowVectorPtr> generateInput(
                const std::vector<std::string> &keyNames,
                const std::vector<TypePtr> &keyTypes);

            void addPlansWithTableScan(
                const std::string &tableDir,
                const std::vector<std::string> &partitionKeys,
                const std::vector<RowVectorPtr> &input,
                std::vector<PlanWithSplits> &altPlans);

            // Makes the query plan with default settings in RowNumberFuzzer and value
            // inputs for both probe and build sides.
            //
            // NOTE: 'input' could either input rows with lazy
            // vectors or flatten ones.
            static PlanWithSplits makeDefaultPlan(
                const std::vector<std::string> &partitionKeys,
                const std::vector<RowVectorPtr> &input);

            static PlanWithSplits makePlanWithTableScan(
                const RowTypePtr &type,
                const std::vector<std::string> &partitionKeys,
                const std::vector<Split> &splits);
        };

        RowNumberFuzzer::RowNumberFuzzer(
            size_t initialSeed,
            std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner)
            : RowNumberFuzzerBase(initialSeed, std::move(referenceQueryRunner)) {
        }

        std::pair<std::vector<std::string>, std::vector<TypePtr> >
        RowNumberFuzzer::generatePartitionKeys() {
            const auto numKeys = randInt(1, 3);
            std::vector<std::string> names;
            std::vector<TypePtr> types;
            for (auto i = 0; i < numKeys; ++i) {
                names.push_back(fmt::format("c{}", i));
                types.push_back(vectorFuzzer_.randType(/*maxDepth=*/1));
            }
            return std::make_pair(names, types);
        }

        std::vector<RowVectorPtr> RowNumberFuzzer::generateInput(
            const std::vector<std::string> &keyNames,
            const std::vector<TypePtr> &keyTypes) {
            std::vector<std::string> names = keyNames;
            std::vector<TypePtr> types = keyTypes;
            // Add up to 3 payload columns.
            const auto numPayload = randInt(0, 3);
            for (auto i = 0; i < numPayload; ++i) {
                names.push_back(fmt::format("c{}", i + keyNames.size()));
                types.push_back(vectorFuzzer_.randType(/*maxDepth=*/2));
            }

            const auto inputType = ROW(std::move(names), std::move(types));
            std::vector<RowVectorPtr> input;
            input.reserve(turbo::get_flag(FLAGS_num_batches));
            for (auto i = 0; i < turbo::get_flag(FLAGS_num_batches); ++i) {
                input.push_back(vectorFuzzer_.fuzzInputRow(inputType));
            }

            return input;
        }

        RowNumberFuzzerBase::PlanWithSplits RowNumberFuzzer::makeDefaultPlan(
            const std::vector<std::string> &partitionKeys,
            const std::vector<RowVectorPtr> &input) {
            auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
            std::vector<std::string> projectFields = partitionKeys;
            projectFields.emplace_back("row_number");
            auto plan = PlanBuilder()
                    .values(input)
                    .rowNumber(partitionKeys)
                    .project(projectFields)
                    .planNode();
            return PlanWithSplits{std::move(plan)};
        }

        RowNumberFuzzerBase::PlanWithSplits RowNumberFuzzer::makePlanWithTableScan(
            const RowTypePtr &type,
            const std::vector<std::string> &partitionKeys,
            const std::vector<Split> &splits) {
            std::vector<std::string> projectFields = partitionKeys;
            projectFields.emplace_back("row_number");

            auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
            core::PlanNodeId scanId;
            auto plan = PlanBuilder(planNodeIdGenerator)
                    .tableScan(type)
                    .rowNumber(partitionKeys)
                    .project(projectFields)
                    .planNode();
            return PlanWithSplits{plan, splits};
        }

        void RowNumberFuzzer::addPlansWithTableScan(
            const std::string &tableDir,
            const std::vector<std::string> &partitionKeys,
            const std::vector<RowVectorPtr> &input,
            std::vector<PlanWithSplits> &altPlans) {
            POLLUX_CHECK(!tableDir.empty());

            if (!isTableScanSupported(input[0]->type())) {
                return;
            }

            const std::vector<Split> inputSplits = test::makeSplits(
                input, fmt::format("{}/row_number", tableDir), writerPool_);
            altPlans.push_back(makePlanWithTableScan(
                as_row_type(input[0]->type()), partitionKeys, inputSplits));
        }

        void RowNumberFuzzer::runSingleIteration() {
            const auto [keyNames, keyTypes] = generatePartitionKeys();

            const auto input = generateInput(keyNames, keyTypes);
            test::logVectors(input);

            auto defaultPlan = makeDefaultPlan(keyNames, input);
            const auto expected =
                    execute(defaultPlan, pool_, /*injectSpill=*/false, false);

            if (expected != nullptr) {
                validateExpectedResults(defaultPlan.plan, input, expected);
            }

            std::vector<PlanWithSplits> altPlans;
            altPlans.push_back(std::move(defaultPlan));

            const auto tableScanDir = exec::test::TempDirectoryPath::create();
            addPlansWithTableScan(tableScanDir->getPath(), keyNames, input, altPlans);

            for (auto i = 0; i < altPlans.size(); ++i) {
                testPlan(
                    altPlans[i], i, expected, "core::QueryConfig::kRowNumberSpillEnabled");
            }
        }
    } // namespace

    void rowNumberFuzzer(
        size_t seed,
        std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner) {
        RowNumberFuzzer(seed, std::move(referenceQueryRunner)).run();
    }
} // namespace kumo::pollux::exec
